//--------------------------------------------------------------------------
// $Id: GenericConsumer.java 12 2006-08-28 05:05:13Z zendragon $
//--------------------------------------------------------------------------
// Copyright 2006 Social Brain Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//--------------------------------------------------------------------------

package org.socialbrain.openid.consumer;

import org.socialbrain.openid.Assertion;
import org.socialbrain.openid.AssociateRequest;
import org.socialbrain.openid.Association;
import org.socialbrain.openid.AuthenticateRequest;
import org.socialbrain.openid.Constants;
import org.socialbrain.openid.KeyValuePairs;
import org.socialbrain.openid.OpenIdException;
import org.socialbrain.openid.Utils;
import org.socialbrain.openid.security.DHCrypto;
import org.socialbrain.openid.security.DHKeyPair;

import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpStatus;
import org.apache.commons.httpclient.methods.PostMethod;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import java.io.IOException;
import java.net.MalformedURLException;
import java.util.List;
import java.util.Map;

/**
 * @author <a href="http://erning.net/">Zhang Erning</a>
 * @version $Id: GenericConsumer.java 12 2006-08-28 05:05:13Z zendragon $
 */
public class GenericConsumer implements Consumer
{
    public ConsumerToken begin(String identifier, String returnTo, String trustRoot)
    {
        return begin(identifier, returnTo, trustRoot, false);
    }

    public ConsumerToken begin(String identifier, String returnTo, String trustRoot, boolean immediate)
    {
        Discovered discovered = discovery(identifier);
        String endpoint = discovered.getEndpoint();
        String identity = StringUtils.isNotEmpty(discovered.getDelegateIdentifier())
                ? discovered.getDelegateIdentifier()
                : discovered.getClaimedIdentifier();

        AssociationSession session = _associationSessionManager.lookupByEndpoint(endpoint);
        if (session == null)
        {
            session = associate(endpoint, _defaultAssocType, _defaultSessionType);
            _associationSessionManager.put(identity, session);
        }
        String redirectURL = authenticate(identity, returnTo, trustRoot, immediate);

        return new ConsumerToken(identity, redirectURL);
    }

    public boolean complete(String identity, String queryString, String encoding)
    {
        KeyValuePairs pairs = KeyValuePairs.parseQueryString(queryString, encoding);
        Assertion assertion = new Assertion(pairs);
        return verifyAuthentication(identity, assertion);
    }

    public boolean complete(String identity, Map<String, String> parameters)
    {
        KeyValuePairs pairs = KeyValuePairs.parseParameters(parameters);
        Assertion assertion = new Assertion(pairs);
        return verifyAuthentication(identity, assertion);
    }

    /**
     * Discovery
     *
     * <a href="http://www.openid.net/specs/openid-authentication-2_0-08.html#discovery">specs</a>
     *
     * @param identifier
     * @return Discovered information or null on fail.
     */
    public Discovered discovery(String identifier)
    {
        identifier = Utils.normalizeIdentifier(identifier);

        for (Discoverer discoverer : _discoverers)
        {
            Discovered response = discoverer.discover(identifier);
            if (response != null)
            {
                return response;
            }
        }

        return null;
    }

    /**
     * Diffie-Hellman association sessions.
     *
     * <a href="http://www.openid.net/specs/openid-authentication-2_0-08.html#dh_sessions">specs</a>
     *
     * @param endpoint IdP's endpoint
     * @param assocType "HMAC-SHA1" or "HMAC-SHA256"
     * @param sessionType "DH-SHA1" or "DH-SHA256"
     * @return AssociationSession object for futher use
     */
    public AssociationSession associate(String endpoint, String assocType, String sessionType)
    {
        AssociateRequest request = new AssociateRequest();
        request.setAssocType(assocType);
        request.setSessionType(sessionType);

        if (!request.isUnencrypted())
        {
            DHCrypto dhCrypto = findDhCrypto(assocType, sessionType);
            if (dhCrypto == null)
            {
                _logger.debug("no supported DHCrypto found. Use unencrypted session.");
            }
            else
            {
                request.setSessionType(sessionType);
                request.setDhModulus(_dhKeyPair.getModulus());
                request.setDhGen(_dhKeyPair.getGen());
                request.setDhConsumerPublic(_dhKeyPair.getPublicKey());
            }
        }

        Association association = associate(endpoint, request);

        if (association.isError())
        {
            // TODO:
            throw new OpenIdException("association.isError");
        }

        if (association.isUnencrypted())
        {
            if (_logger.isDebugEnabled() && !request.isUnencrypted())
            {
                _logger.debug("Server does not support encryption type " + sessionType);
            }
            return new AssociationSession(endpoint,
                    association.getAssocHandle(), association.getAssocType(), association.getExpireIn(),
                    association.getMacKey());
        }

        // TODO: will server returns different assocType and/or sessionType?
        DHCrypto dhCrypto = findDhCrypto(association.getAssocType(), association.getSessionType());
        if (dhCrypto == null)
        {
            // TODO: unsupported assocType and/or sessionType server wanted?
            throw new OpenIdException("no supported DHCrypto found");
        }

        return new AssociationSession(endpoint,
                association.getAssocHandle(), association.getAssocType(), association.getExpireIn(),
                dhCrypto.decrypt(association.getEncMacKey(), _dhKeyPair, association.getDhServerPublic()));
    }

    protected DHCrypto findDhCrypto(String assocType, String sessionType)
    {
        for (DHCrypto dhCrypto : _dhCryptos)
        {
            if (dhCrypto.isAcceptAssocType(assocType) && dhCrypto.isAcceptSessionType(sessionType))
            {
                return dhCrypto;
            }
        }
        return null;
    }

    /**
     *
     * @param endpoint
     * @param request
     * @return Association
     */
    protected Association associate(String endpoint, AssociateRequest request)
    {
        PostMethod method = new PostMethod(endpoint);
        for (String key : request.keySet())
        {
            method.addParameter(Constants.PARAMETER_PREFIX + key, request.get(key));
        }
        try
        {
            if (_httpClient.executeMethod(method) != HttpStatus.SC_OK)
            {
                throw new OpenIdException(String.format("Unexcept http status: %03d - %s",
                        method.getStatusCode(), method.getStatusText()));
            }

            String formString = Utils.readStringFromStream(method.getResponseBodyAsStream(),
                    method.getResponseContentLength(), method.getResponseCharSet());

            KeyValuePairs pairs = KeyValuePairs.parseFormString(formString);
            return new Association(pairs);
        }
        catch (IOException e)
        {
            throw new OpenIdException(
                    String.format("Unable to associate with openid server at %s", endpoint), e);
        }
        finally
        {
            method.releaseConnection();
        }
    }

    /**
     * Returns the URL for checkid_immediate or checkid_setup
     *
     * @param identity
     * @param returnTo
     * @param trustRoot
     * @param immediate
     * @return The URL the user agent will redirect to
     */
    public String authenticate(String identity, String returnTo, String trustRoot, boolean immediate)
    {
        AssociationSession session = _associationSessionManager.lookupByIdentity(identity);
        if (session == null)
        {
            throw new OpenIdException(String.format("Session for %s not found", identity));
        }

        String endpoint = session.getEndpoint();
        String assocHandle = session.getAssocHandle();

        AuthenticateRequest request = new AuthenticateRequest(immediate);
        if (identity != null) request.setIdentity(identity);
        if (assocHandle != null) request.setAssocHandle(assocHandle);
        request.setReturnTo(returnTo);
        request.setTrustRoot(trustRoot != null ? trustRoot : returnTo);
        request.setAssocHandle(assocHandle);

        try
        {
            return Utils.appendQueryString(endpoint, request.toQueryString());
        }
        catch (MalformedURLException e)
        {
            throw new OpenIdException("Malformed endpoint", e);
        }
    }

    /**
     * Verify the authentication by the query string returned by IdP.
     *
     * @param identity identity to be verified
     * @param assertion Assertion to be verified
     * @return authenticated or not
     */
    protected boolean verifyAuthentication(String identity, Assertion assertion)
    {
        for (Verifier verifier : _verifiers)
        {
            if (!verifier.verify(identity, assertion))
            {
                return false;
            }
        }
        return true;
    }

    //

    public void setAssociationSessionManager(AssociationSessionManager value)
    {
        _associationSessionManager = value;
    }

    public AssociationSessionManager getAssociationSessionManager()
    {
        return _associationSessionManager;
    }

    private AssociationSessionManager _associationSessionManager;

    //

    public void setDiscoverers(List<Discoverer> value)
    {
        _discoverers = value;
    }

    public List<Discoverer> getDiscoverers()
    {
        return _discoverers;
    }

    private List<Discoverer> _discoverers;

    //

    public void setDhKeyPair(DHKeyPair value)
    {
        _dhKeyPair = value;
    }

    public DHKeyPair getDHKeyPair()
    {
        return _dhKeyPair;
    }

    public void setDhCryptos(List<DHCrypto> value)
    {
        _dhCryptos = value;
    }

    public List<DHCrypto> getDhCryptos()
    {
        return _dhCryptos;
    }

    private DHKeyPair _dhKeyPair;
    private List<DHCrypto> _dhCryptos;

    //

    public void setDefaultAssocType(String value)
    {
        _defaultAssocType = value;
    }

    public String getDefaultAssocType()
    {
        return _defaultAssocType;
    }

    public void setDefaultSessionType(String value)
    {
        _defaultSessionType = value;
    }

    public String getDefaultSessionType()
    {
        return _defaultSessionType;
    }

    private String _defaultAssocType = Constants.DEFAULT_ASSOC_TYPE;
    private String _defaultSessionType = Constants.DEFAULT_SESSION_TYPE;

    //

    public void setVerifiers(List<Verifier> value)
    {
        _verifiers = value;
    }

    public List<Verifier> getVerifiers()
    {
        return _verifiers;
    }

    private List<Verifier> _verifiers;

    //

    public void setHttpClient(HttpClient value)
    {
        _httpClient = value;
    }

    public HttpClient getHttpClient()
    {
        return _httpClient;
    }

    private HttpClient _httpClient;

    private transient static final Log _logger = LogFactory.getLog(GenericConsumer.class);
}
